[Pytorch][Common] Hybrid quantization#2817
Conversation
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR introduces hybrid (per-direction) quantization for PyTorch: a
Confidence Score: 4/5Safe to merge for non-Hopper FSDP+distopt workflows; the Hopper/L40 columnwise-only hybrid Float8 distopt path in _cast_master_weights_to_fp8_current_scaling will crash for affected configurations. The FSDP non-shard path for Hopper columnwise-only hybrid Float8 distopt has been protected by _update_transpose_only_float8_flat_fragment, but the shard path (use_fsdp_shard_model_weights=True) calls quantizer.update_quantized directly on a Float8Tensor with _data=None, crashing the C++ kernel. This affects Hopper/L40 + distributed optimizer + hybrid Float8 current-scaling — a combination the PR description explicitly marks as done. transformer_engine/pytorch/tensor/utils.py (_cast_master_weights_to_fp8_current_scaling, lines 561–574): the FSDP shard branch needs the same _update_transpose_only_float8_flat_fragment guard as the non-FSDP branch for Hopper columnwise-only Float8 sub-storages. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
HQ[HybridQuantizer\nrowwise_quantizer\ncolumnwise_quantizer\ncolumnwise_source] --> |quantize_impl| HT[HybridQuantizedTensor\nrowwise_storage\ncolumnwise_storage]
HT --> |_unwrap_hybrid_A/B| RW[RowwiseTensorStorage\nFloat8/MXFP8/NVFP4/Identity]
HT --> |_unwrap_hybrid_A/B| CW[ColumnwiseTensorStorage\nFloat8/MXFP8/Identity]
RW --> |_unwrap_identity_tensor| GEMM[general_gemm / grouped_gemm\nC++ kernel]
CW --> |_unwrap_identity_tensor| GEMM
HT --> |fsdp_pre_all_gather| FSDP[FSDP2 all-gather\nrow_buffers + col_buffers]
FSDP --> |fsdp_post_all_gather| HT2[HybridQuantizedTensor\nfull-size param]
HQ --> |quantize_master_weights| ROUTE[_route_hybrid_to_buckets\nper-direction dispatch]
ROUTE --> |Float8Quantizer| DSP[delayed_scaling_params]
ROUTE --> |Float8CurrentScalingQuantizer| CSP[current_scaling_params]
ROUTE --> |IdentityQuantizer| IDP[identity_params]
CSP --> |non-FSDP: guarded| CAST1[_update_transpose_only\n_float8_flat_fragment OK]
CSP --> |FSDP: unguarded| CRASH[quantizer.update_quantized\n_data=None crash on Hopper]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
HQ[HybridQuantizer\nrowwise_quantizer\ncolumnwise_quantizer\ncolumnwise_source] --> |quantize_impl| HT[HybridQuantizedTensor\nrowwise_storage\ncolumnwise_storage]
HT --> |_unwrap_hybrid_A/B| RW[RowwiseTensorStorage\nFloat8/MXFP8/NVFP4/Identity]
HT --> |_unwrap_hybrid_A/B| CW[ColumnwiseTensorStorage\nFloat8/MXFP8/Identity]
RW --> |_unwrap_identity_tensor| GEMM[general_gemm / grouped_gemm\nC++ kernel]
CW --> |_unwrap_identity_tensor| GEMM
HT --> |fsdp_pre_all_gather| FSDP[FSDP2 all-gather\nrow_buffers + col_buffers]
FSDP --> |fsdp_post_all_gather| HT2[HybridQuantizedTensor\nfull-size param]
HQ --> |quantize_master_weights| ROUTE[_route_hybrid_to_buckets\nper-direction dispatch]
ROUTE --> |Float8Quantizer| DSP[delayed_scaling_params]
ROUTE --> |Float8CurrentScalingQuantizer| CSP[current_scaling_params]
ROUTE --> |IdentityQuantizer| IDP[identity_params]
CSP --> |non-FSDP: guarded| CAST1[_update_transpose_only\n_float8_flat_fragment OK]
CSP --> |FSDP: unguarded| CRASH[quantizer.update_quantized\n_data=None crash on Hopper]
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
timmoon10
left a comment
timmoon10
left a comment
There was a problem hiding this comment.
Overall I think this moves us in a good direction. I see some minor bugs, as well as bugs reported by @greptile-apps.
| rowwise_result = self.rowwise_quantizer.quantize(tensor) | ||
| columnwise_result = self.columnwise_quantizer.quantize(tensor) |
There was a problem hiding this comment.
Do we handle the case where not all usages are needed? I'd expect something like:
| rowwise_result = self.rowwise_quantizer.quantize(tensor) | |
| columnwise_result = self.columnwise_quantizer.quantize(tensor) | |
| rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None | |
| columnwise_result = self.columnwise_quantizer.quantize(tensor) if self.columnwise_usage else None |
| requires_grad: bool = False, | ||
| pin_memory: bool = False, | ||
| ) -> HybridQuantizedTensor: | ||
| self.rowwise_quantizer.internal = True |
There was a problem hiding this comment.
Could we just set internal=True in the constructor? I don't think we ever need PyTorch tensor functionality in the per-usage data.
There was a problem hiding this comment.
This would not work under FSDP2.
| def factory(role): | ||
| if role == "linear_weight": | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_fp8_quantizer(), | ||
| columnwise_quantizer=_make_mxfp8_quantizer(), | ||
| ) | ||
| if role == "linear_input": | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_fp8_quantizer(), | ||
| columnwise_quantizer=_make_nvfp4_quantizer(), | ||
| ) | ||
| if role in ("linear_grad_output", "linear_grad_input"): | ||
| return HybridQuantizer( | ||
| rowwise_quantizer=_make_mxfp8_quantizer(), | ||
| columnwise_quantizer=_make_nvfp4_quantizer(), | ||
| ) | ||
| return None |
There was a problem hiding this comment.
This is horrifying. Good test.
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
| # DCP serializes ``CustomRecipe`` via ``pickle``; closure-based qfactories | ||
| # (lambdas, inner functions referencing captured state) are not picklable, | ||
| # so the qfactory must live at module scope. See | ||
| # ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. |
There was a problem hiding this comment.
This comment is potentially useful, but I don't think it is in the right place - shouldn't it be closer to the actual implementation?
| for param in model.parameters(): | ||
| state = optimizer.state[param] | ||
| assert state["exp_avg"].dtype == torch.float32 | ||
| assert state["exp_avg_sq"].dtype == torch.float32 | ||
| if "master_param" in state: | ||
| assert state["master_param"].dtype == torch.float32 | ||
|
|
||
| assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" |
There was a problem hiding this comment.
That's not a very strict test, is there a way for us to do some numerical correctness comparisons?
There was a problem hiding this comment.
Enabled check for the monotonic loss decrease (still mostly sanity), and also enabled hybrid vs vanilla bitwise recipe comparizon, see e.g. test_fused_adam_hybrid_vs_base_recipe_parity.
Signed-off-by: Evgeny <etsykunov@nvidia.com>
kwyss-nvidia
left a comment
kwyss-nvidia
left a comment
There was a problem hiding this comment.
Thanks Evgeny for this expansive PR!
I'm excited to see the columnwise_source options in hybrid quantizer and that the edge cases for the FSDP protocol are considered and captured in the new tensor types.
LGTM!
| assert copied.rowwise_usage is False | ||
| assert copied.columnwise_usage is True | ||
|
|
||
| def test_rowwise_dequantized_identity_columnwise_matches_rowwise(self, input_tensor): |
There was a problem hiding this comment.
Thanks for adding the coverage for double quantization and the field so that the quantized tensor tracks describes columnwise source.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| # Module-level qfactories (picklable, required for checkpoint serialization). |
There was a problem hiding this comment.
Which qfactories go into the checkpoint serialization?
While it's useful to have provenance of how the checkpoint was created, does the pickling of qfactories mean that the resulting checkpoints won't be read by transformer engine's with the same classes for custom quantization.
Is there any way to override this requirement and load the checkpoint as BF16, ignoring the pickled qfactories?
There was a problem hiding this comment.
Good point. I clarified the comment: the module-level qfactory is only needed so TE-to-TE quantized-param checkpoints have a stable importable reference for any pickled quantizer/recipe metadata.
For portability: if quantized_model_init is disabled, model weights are normal BF16 tensors and the CustomRecipe extra state is not needed for stateless recipes, so an external consumer can ignore TE _extra_state. With quantized_model_init, the model state_dict stores TE quantized tensor subclasses for any recipe, not just hybrid/CustomRecipe, so TE -> third-party runtime portability would need a separate high-precision/BF16 export path for quantized primary weights.
Would it be useful to plan that quantized_model_init BF16-weight state_dict/export support as a follow-up, or is running without quantized_model_init for portable checkpoints sufficient for your use case?
There was a problem hiding this comment.
Running without quantized_model_init is sufficient. I did not know about that option!
| transa = layout[0] == "T" | ||
| transb = layout[1] == "T" | ||
|
|
||
| A = _materialize_high_precision(_unwrap_hybrid_A(A, layout)) |
There was a problem hiding this comment.
The naming convention could be revisited. It can be interpreted easily but falsely to mean all quantized tensors will be materialized into high precision tensors. Perhaps "_unwrap_if_high_precision"?
There was a problem hiding this comment.
Replaced with _unwrap_identity_tensor, since we are doing isinstance(tensor, IdentityTensorStorage)
|
|
||
| # Linear-only recipe (no attention quantization): the qfactory is the only knob. | ||
| recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) | ||
| with autocast(recipe=recipe): |
There was a problem hiding this comment.
This is pleasantly simple as an API. Thanks Evgeny.
| ``HybridQuantizer`` terms, that source choice is expressed with | ||
| ``columnwise_source="rowwise_dequantized"``. | ||
|
|
||
| All non-weight roles keep the standard NVFP4 factory behavior, including RHT |
There was a problem hiding this comment.
This is useful. We have trained with an equivalent recipe for several experiments and I'm looking forward to trying this implementation.
| # Return early if recipe state matches recipe | ||
| if self.fp8_meta_tensors_initialized: | ||
| recipe_state = self.fp8_meta[fp8_meta_tensor_key] | ||
| # Follow-up: Match built-in recipes by full config, not just RecipeState type, so |
There was a problem hiding this comment.
If this follow up liked an issue or this pull request ID, it would be easier to grep for all related follow ups.
There was a problem hiding this comment.
Create an umbrella tracker #3158 and referenced it
| ------- | ||
| MXFP8 forward plus high-precision backward from the rowwise-dequantized | ||
| forward value can be expressed as:: | ||
|
|
There was a problem hiding this comment.
In the zoo, there's an example where the weight tensor is specialized with double quantization. It would be helpful to illustrate that it's possible to customize along the weight/activation/grad axis as well as the rowwise/colwise abstraction and reference the zoo.
| for sub in (self.rowwise_quantizer, self.columnwise_quantizer): | ||
| group = getattr(sub, "amax_reduction_group", None) | ||
| if group is not None: | ||
| return group |
There was a problem hiding this comment.
Arguably, this should assert if there are two groups, they are consistent. Is that possible?
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: root <root@prenyx0286.a51.clusters.nvidia.com>
| # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP | ||
| # activations are gathered in high precision and re-quantized whole, so | ||
| # every rank already sees the same global amax. | ||
| # TODO(#3158): once native quantized all-gather lands (see |
There was a problem hiding this comment.
Do you intend to close those TODOs?
There was a problem hiding this comment.
Yes, but in the follow-up PRs, see #3158. for now it falls back to high-precision AG
| if ( | ||
| save_original_input | ||
| and backward_needs_input | ||
| and input_quantizer is not None | ||
| and not input_quantizer.allows_save_original_input_for_backward() | ||
| ): | ||
| warnings.warn( | ||
| "Ignoring save_original_input=True because the input quantizer requires " | ||
| "the forward quantized activation for backward " | ||
| f"({input_quantizer}).", | ||
| stacklevel=2, | ||
| ) | ||
| save_original_input = False |
There was a problem hiding this comment.
Could this check be moved to the quantizer contruction rather than being called on every forward pass?
save_original_input is dependent on module options and recipe.
There was a problem hiding this comment.
I would keep this check at runtime for now. save_original_input depends on thiings that might be changed at runtime right?
| qkv_quantizer = recipe.qfactory( | ||
| QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) | ||
| ) |
There was a problem hiding this comment.
This does not look safe (and also is wasteful) to call this every time. E.g. if the custom quantizer
is stateful or uses RNG etc. we would be redoing seeding of it every time potentially changing the
result depending on how many forwards of MHA we called.
There was a problem hiding this comment.
Fixed. The QKV quantizer is now materialized once through DPA recipe state and reused. MHA no longer calls qfactory to probe it.
| """ | ||
|
|
||
| _hp_data: Optional[torch.Tensor] | ||
| _quantizer: Optional[Quantizer] |
There was a problem hiding this comment.
_quantizer is already a field in the base class.
| """Return the held high-precision tensor (no-op dequantization).""" | ||
| if self._hp_data is None: | ||
| raise RuntimeError("IdentityTensorStorage has no data to dequantize") | ||
| if dtype is not None and self._hp_data.dtype != dtype: |
There was a problem hiding this comment.
Fixed. It now defaults to its nominal dtype
| # Return early if recipe state matches recipe | ||
| if self.fp8_meta_tensors_initialized: | ||
| recipe_state = self.fp8_meta[fp8_meta_tensor_key] | ||
| # Follow-up (#3157): Match built-in recipes by full config, not just RecipeState type, so |
There was a problem hiding this comment.
Can we change "Follow-up" to "TODO"? It is more discoverable that way.
| return False | ||
|
|
||
|
|
||
| def _has_identity_quantizer_list(quantizers): |
There was a problem hiding this comment.
This should not be called every time and instead we should only update this when the recipe changes.
| non_none = [q for q in quantizers if q is not None] | ||
| if not non_none: | ||
| return False | ||
| hybrid_count = sum(1 for q in non_none if isinstance(q, HybridQuantizer)) | ||
| if hybrid_count == 0: | ||
| return False |
There was a problem hiding this comment.
This is pretty convoluted and inefficient.
| """ | ||
| from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage | ||
|
|
||
| if not all(isinstance(q, HybridQuantizer) for q in quantizers): |
There was a problem hiding this comment.
But you already check that in the function above, why checking it again?
| from ...debug.pytorch.debug_quantization import DebugQuantizer | ||
| from ...debug.pytorch.debug_state import TEDebugState | ||
|
|
||
|
|
There was a problem hiding this comment.
A general comment to this entire file: this introduces quite a lot of CPU overhead to this module.
We should just limit this to only allow the same quantizers everywhere rather than checking every
time.
There was a problem hiding this comment.
General comment for the grouped_linear.py
This and the three related comments above are valid and resolved. See c7e11d2
| if out_data is not None: | ||
| out_shape = out_data.size() | ||
| else: | ||
| out_shape = torch.empty(tuple(self.size()), device="meta").view(shape).shape |
There was a problem hiding this comment.
Huh. That's pretty roundabout way of calculating that...
There was a problem hiding this comment.
Right, take a look at the fix at c7e11d2
| """The sub-storage providing columnwise quantized data.""" | ||
| return self._columnwise_storage | ||
|
|
||
| def update_usage( |
There was a problem hiding this comment.
Asking for a storage that does not exist should result in error.
There was a problem hiding this comment.
Added a check for missing directions
| def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: | ||
| """Resolve PyTorch view shape syntax, including ``-1`` inference.""" | ||
| return torch.empty(tuple(input_shape), device="meta").view(*shape).shape |
There was a problem hiding this comment.
That's the same formula that was used in one of the other files. If useful, it should be somewhere where it could be used without rediscovering, but still I think it is pretty roundabout...
There was a problem hiding this comment.
Additionally using torch.empty, view etc would introduce quite a bit of CPU overheads.. We can define a function to do manual size resolution in python that resolves -1, 0 etc.
Here something Gemini gave me
def get_actual_view_shape_full(input_shape: tuple, view_shape: tuple) -> tuple:
# 1. First, resolve any '0' dimensions by copying from the input_shape
resolved_zeros = []
for i, dim in enumerate(view_shape):
if dim == 0:
if i >= len(input_shape):
raise ValueError("Dimension '0' exceeds input shape dimensions")
resolved_zeros.append(input_shape[i])
else:
resolved_zeros.append(dim)
# 2. Calculate element totals
total_elements = math.prod(input_shape)
known_elements = math.prod(d for d in resolved_zeros if d != -1)
# 3. Resolve the '-1' dimension if it exists
if -1 in resolved_zeros:
if resolved_zeros.count(-1) > 1:
raise ValueError("Only one dimension can be -1")
if total_elements % known_elements != 0:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
missing_dim = total_elements // known_elements
return tuple(missing_dim if d == -1 else d for d in resolved_zeros)
# 4. Strict validation if no -1 is used
if total_elements != known_elements:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
return tuple(resolved_zeros)
There was a problem hiding this comment.
Now float8 tensor and storage share the same better shaped utility, see c7e11d2
| ) | ||
|
|
||
| @classmethod | ||
| def __torch_dispatch__(cls, func, types, args, kwargs=None): |
There was a problem hiding this comment.
Do we need this function to be this big? We should be able to mostly just delegate to the held tensor, right?
There was a problem hiding this comment.
Good point, delegated
| "instance or a QuantizerRequest. For an intentional no-op " | ||
| "(high-precision / unquantized) slot, return an " | ||
| "IdentityQuantizer instead of None." |
There was a problem hiding this comment.
We could probably just allow None to mean the IdentityQuantizer here?
There was a problem hiding this comment.
I would not do that. Current behavior is explicit. Treating a factory return value of None as an Identity quantizer would silently broaden the custom recipe API and could hide a missing role. Also none quantizer has a semantic meaning in the modules, which is different from the IdentityQuantizer, right? I suggest Identity should be requested explicitly.
| def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: | ||
| """Resolve PyTorch view shape syntax, including ``-1`` inference.""" | ||
| return torch.empty(tuple(input_shape), device="meta").view(*shape).shape |
There was a problem hiding this comment.
Additionally using torch.empty, view etc would introduce quite a bit of CPU overheads.. We can define a function to do manual size resolution in python that resolves -1, 0 etc.
Here something Gemini gave me
def get_actual_view_shape_full(input_shape: tuple, view_shape: tuple) -> tuple:
# 1. First, resolve any '0' dimensions by copying from the input_shape
resolved_zeros = []
for i, dim in enumerate(view_shape):
if dim == 0:
if i >= len(input_shape):
raise ValueError("Dimension '0' exceeds input shape dimensions")
resolved_zeros.append(input_shape[i])
else:
resolved_zeros.append(dim)
# 2. Calculate element totals
total_elements = math.prod(input_shape)
known_elements = math.prod(d for d in resolved_zeros if d != -1)
# 3. Resolve the '-1' dimension if it exists
if -1 in resolved_zeros:
if resolved_zeros.count(-1) > 1:
raise ValueError("Only one dimension can be -1")
if total_elements % known_elements != 0:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
missing_dim = total_elements // known_elements
return tuple(missing_dim if d == -1 else d for d in resolved_zeros)
# 4. Strict validation if no -1 is used
if total_elements != known_elements:
raise ValueError(f"Cannot reshape {input_shape} to {view_shape}")
return tuple(resolved_zeros)
| raise NotImplementedError( | ||
| "FSDP2 is only supported for MXFP8Tensors with compact scales" | ||
| ) | ||
| names = self.fsdp_buffer_fields() |
There was a problem hiding this comment.
I think it would be good to reuse fsdp extract buffers and fsdp_assign_gather for the current tensor class as well along with hybrid tensor(in this case mxfp8 and similarly for others)
| return dst | ||
|
|
||
| # ── FSDP2: new_zeros ───────────────────────────────────────── | ||
| if func == aten.new_zeros.default: |
There was a problem hiding this comment.
This is risky if we use it for anything apart from FSDP2. Might make sense to zero out the substorages.
| # Factories stay small; these tests target TP/SP plumbing. | ||
|
|
||
|
|
||
| def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): |
There was a problem hiding this comment.
We now recommend to use transformer_engine.pytorch.DType instead of tex.DType in TE. So we should change it here as well
Signed-off-by: Evgeny <etsykunov@nvidia.com>
…ation in GroupedLinear, etc Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
|
/te-ci pytorch L1 |
for more information, see https://pre-commit.ci
5 participants
Description
Hybrid (per-direction) quantization. Hybrid means rowwise/colwise can use different formats via CustomRecipe(qfactory).
This is an experimental feature.
The main problem that it tries to solve is that precision requirements are non-uniform.
Current recipes set one format for both rowwise and colwise directions.
Hybrid quantization enables, e.g. MXFP8 fwd and NVFP4 bwd (or vice versa) or any other valid combination. No need for a hardcoded recipe for every combination.
Composer-style (Composer 2 paper) grouped GEMM recipe, e.g. row-scaled NVFP4 fwd + MXFP8 bwd:
By default, the above factory uses
columnwise_source="original", so MXFP8 backward operands are quantized from the original high-precision tensor. Usecolumnwise_source="rowwise_dequantized"when the backward operand should be derived from the dequantized rowwise NVFP4 forward value.C++ optimizations (fusions, etc.) will come as standalone PRs. cc @kainzhong
TODO:
Follow-up issue tracker #3158.
Integration
Ecosystem integration (all functional, unit-tested):
Megatron-LM integration status:
--fp{4,8}-param-gather+ dist opt (persistent low-precision params viaquantized_model_init+ sharded-master FP32 → quantized cast viaquantize_master_weights.)- [Done] Per-tensor Float8 hybrid (delayed and/or current, any per-direction combination
including same-format, cross-format Float8, single-direction)
- [TODO] Per-block hybrid sub-quantizers (MXFP8, NVFP4, Float8Blockwise) — each rejected per-direction by
quantize_master_weights; unblocker is TE-side cast-helper / kernel.--fp{4,8}-param-gather(fix private attribute access)--fp{4,8}-param-gather- [Done] TE-side hybrid FSDP2 path works end-to-end for Float8 / MXFP8 / Float8Blockwise sub-storages (TODO: need some minor MLM update)
- [TODO] NVFP4 sub-storage FSDP2 hooks
_hybrid_split_quantizeunder Megatron MoE)Review
Total diff +14000
New hybrid source (
hybrid_tensor.py,hybrid_tensor_storage.py,identity_tensor.py,identity_tensor_storage.py) ~1800Adjacent modifications ~1500
Tests are the rest (~10K)
Suggested reading order
-columnwise_source controls whether columnwise quantization uses the original input or the rowwise-dequantized value.
1.1 Identity passthrough — b99277a
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: